Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Introduction of model_typed and model_warntype in DebugUtils #708

Open
wants to merge 14 commits into
base: master
Choose a base branch
from

Conversation

torfjelde
Copy link
Member

I wanted to do some checking of type-stability for a model. Unfortunately, this is quite annoying to do in a way where you can inspect the inferred types of the model's evaluator function.

Hence I was thinking it might be useful to have a few simple methods which just does the set up + calls @code_warntype on model.f (and similarly for @code_typed). This PR therefore introduces model_warntype and model_typed in the DebugUtils module which does exactly this.

I think this is particularly useful for users who wants to check type-stability of their model, which IMO should be able to do so without extensive knowledge of DynamicPPL internals.

The downside: have to add InteractiveUtils.jl as a direct dep (though it's part of base, so should be fine).

Thoughts?

torfjelde and others added 2 commits October 31, 2024 12:01
`model_warntype` for the checking of the model's evaluator
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Copy link

codecov bot commented Oct 31, 2024

Codecov Report

Attention: Patch coverage is 0% with 10 lines in your changes missing coverage. Please review.

Project coverage is 77.67%. Comparing base (b96e368) to head (2925ee1).

Files with missing lines Patch % Lines
src/debug_utils.jl 0.00% 10 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##           master     #708      +/-   ##
==========================================
- Coverage   82.14%   77.67%   -4.47%     
==========================================
  Files          30       30              
  Lines        4200     3919     -281     
==========================================
- Hits         3450     3044     -406     
- Misses        750      875     +125     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

@coveralls
Copy link

coveralls commented Oct 31, 2024

Pull Request Test Coverage Report for Build 11684382988

Details

  • 0 of 10 (0.0%) changed or added relevant lines in 1 file are covered.
  • No unchanged relevant lines lost coverage.
  • Overall coverage decreased (-0.2%) to 82.182%

Changes Missing Coverage Covered Lines Changed/Added Lines %
src/debug_utils.jl 0 10 0.0%
Totals Coverage Status
Change from base Build 11667882952: -0.2%
Covered Lines: 3450
Relevant Lines: 4198

💛 - Coveralls

Copy link
Member

@yebai yebai left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it is very helpful. Also happy to depend on InteractiveUtils.

@willtebbutt you might want to help review this PR too.

Copy link
Member

@willtebbutt willtebbutt left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I really like this, and agree that InteractiveUtils is a perfectly fine dep to have.

So, I think that the concept is basically fine. Will this stuff automatically appear in the docs, or do we need to add it in manually?

src/debug_utils.jl Show resolved Hide resolved
src/debug_utils.jl Show resolved Hide resolved
@willtebbutt
Copy link
Member

willtebbutt commented Nov 4, 2024

While I think you should feel free to merge this PR without this, could I ask whether you would consider exposing a function which returns the function + arguments that you pass to @code_warntype? i.e. something like

fargs, kwargs = get_args_for_analysis(model, varinfo, context))

with the same default values for varinfo and context as for model_warntype and model_typed?

I guess the implementation would be something like

function get_args_for_analysis(model, varinfo=default_thing, context=default_thing)
    args, kargs = DynamicPPL.make_evaluate_args_and_kwargs(model, varinfo, context)
    return (f, args...), kwargs
end

The reason to want this, is that

  1. sometimes Base.code_ircode / Base.code_ircode_by_type is the thing you're after, rather than Base.code_typed (the former returns IRCode while the latter returned CodeInfo), and
  2. if I want to inspect the IR generated by Mooncake.jl for a particular Turing.jl model (which I often want to do in order to debug stuff), I need to have access to the signature that I'm building the rule for, and I can get this from the arguments that you pass into @code_warntype.

@torfjelde
Copy link
Member Author

Yep, can do @willtebbutt !

@torfjelde
Copy link
Member Author

Wait; why is that function you defined simpler than just doing DynamicPPL.make_evaluate_args_and_kwargs(model, varinfo, context)?

@torfjelde
Copy link
Member Author

Added optimize kwarg to the methods 👍

torfjelde and others added 2 commits November 4, 2024 17:15
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
@willtebbutt
Copy link
Member

willtebbutt commented Nov 4, 2024

Wait; why is that function you defined simpler than just doing DynamicPPL.make_evaluate_args_and_kwargs(model, varinfo, context)?

It's a simpler interface, rather than implementation.

It's about ensuring that we have a single source of truth that you can point people towards (and put in the docs) that says "here is everything you need -- we promise that this is the thing that you should profile for performance", as opposed to having to know that the f field of model is the right function to look at (which I assume is an implementation detail rather than a guarantee that the interface makes), and that it must be combined with the output of make_evaluate_arg_and_kwargs.

edit: having now looked at this a bit more closely, am I correct that you can always do:

fargs, kwargs = DynamicPPL.make_evaluate_args_and_kwargs(model, varinfo, context)
fargs[1](fargs[2:end]...; kwargs...)

in order to evaluate the correct thing?

@willtebbutt
Copy link
Member

Also, I've just realised that this functionality is not tested, and doesn't appear in the docs. It would be good to add a unit test which just runs model_typed and model_warntype in order to check that they do indeed run.

Also, could you please add these to a debug section in the docs?

end

"""
model_typed(model[, varinfo, context]; optimize=true)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same here -- needs to specify what the defaults are.

@torfjelde
Copy link
Member Author

: having now looked at this a bit more closely, am I correct that you can always do:

Nah, it's

model.f(args...; kwargs...)

The evaluator f isn't returned by make_evaluator_args_and_kwargs. I guess returning the evaluator might be useful, but is it clear to have it as part of the "args" as you mentioned? Seems a bit confusing to me; why not return (f, args, kwargs) instead of (f, args...), kwargs?

@torfjelde torfjelde self-assigned this Nov 5, 2024
@torfjelde
Copy link
Member Author

Added docs and testing @willtebbutt

@willtebbutt
Copy link
Member

I guess returning the evaluator might be useful, but is it clear to have it as part of the "args" as you mentioned? Seems a bit confusing to me; why not return (f, args, kwargs) instead of (f, args...), kwargs?

I don't feel especially strongly either way. Personally, I like having it with the arguments because for Mooncake.jl-related stuff, where the distinction between function and argument is not really there (e.g. when using Base.code_ircode_by_type you just create the signature for the call). I can also see that most users probably do think about the function as being a separate thing from the args though, so I'm really very happy either way.

Copy link
Member

@willtebbutt willtebbutt left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is ready to go once the patch version has been bumped.

As I said, I don't want my other request to get in the way of getting this released. I'm happy to re-review if you do want to implement that as part of this PR though @torfjelde . I'll open an issue about it if you don't fancy doing it as part of this PR.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants